import os

import jittor
from ljp.dataset import HRRP
from ljp.metrics import accuracy
from ljp import Trainer
from model import Cnn1D

if jittor.has_cuda:
    jittor.flags.use_cuda = 1

batch_size = 64
model = Trainer(Cnn1D(layer_cfg='4l064', kernel_size=7, operator='subtract'))
# model = Trainer(Cnn1D(layer_cfg='4l064', kernel_size=7, operator='mutiply'))
train_loader = HRRP(dset_name='train', batch_size=batch_size, shuffle=True)
val_loader = HRRP(dset_name='test', batch_size=batch_size, shuffle=True)
model.compile(loss_dict={
    'loss_func': jittor.nn.CrossEntropyLoss
},
    optimizer_dict={
        'optimizer': jittor.optim.Adan,
        'lr': 1e-3
    },
    metrics_dict={"accuracy": accuracy},
    monitor='val_accuracy',
    monitor_mode='max'
)
model.fit(epochs=200, T_loader=train_loader, V_loader=val_loader)
